"""
A class to exercise the JT core functions.
"""
import time
import pylab
import numpy as np

import itk
import vtk

import connect2
import ImageCompare

import JT
import Calibration

def Test():
    """
    Test out the ITK registration procedure.
    """
#===============================================================================
#    Data Paths
#===============================================================================
    inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
    fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
    fixedImageMaskName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ucharb2Rmask080-LAT.mha'
    maskFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1b2VOI.mha'
    staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
    calFile = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'
    startPositionsPath = 'C:/Users/bryan/bryan-code/2D3D/startpositions_200_ordered.txt'
    
#===============================================================================
#    Generate Lateral Projection
#===============================================================================
    xraycam = JT.XRayCamera()
    drr = JT.DRR()
    
    drr.SetXRayCamera(xraycam)
    drr.SetVolumeFilename(inputVolFileName)
    drr.SetBlackOnWhiteImage()
    drr.InteractionOff()
    
    # Set Fixed image using path
    fixedImage = JT.FixedImage()
    fixedImage.SetFileName(fixedImageName)
    
    reg=JT.Registration()
    reg.SetFixedImage(fixedImage.GetImage(0))
    reg.SetMovingImageGenerator(drr)
    
    #Set Mask
    #drr._volume.SetImageMask(maskFileName)
    
#===============================================================================
#    Calibration
#===============================================================================
    cal = Calibration.ExternalCalibration()
    cal.LoadConsolidatedCalibration(calFile)
    cal.LoadStaFile(staFile)
    
    xraycam.SetExternalCalibration(cal)
    
    drr._volume.SetOriginalTransform(cal._VolumeT)
    drr._volume.SetBodyCenter((3.906, 2.17, -7.378))
    
#===============================================================================
#    Setup Registration
#===============================================================================
    # Set Fixed image using path
    fixedImage = JT.FixedImage()
    fixedImage.SetFileName(fixedImageName)
    
    reg=JT.Registration()
    reg.SetFixedImage(fixedImage.GetImage(0))
    reg.SetMovingImageGenerator(drr)
    reg._blur = 0.0
    
    #Set Mask (choose one, moving or fixed)
    #drr._volume.SetImageMask(maskFileName)
    #reg.SetFixedImageMask(fixedImageMaskName)
    
    #Based on the fluoro mask, i determined the following region of interest 
    # (row,column) format for numpy array:
    roi = [[210,187],[277,370]]
    reg.SetRegionOfInterest(roi)
    reg.SetImageMetric('gc')


#===============================================================================
#     Adjust Display Settings
#===============================================================================
    rstart = 0
    rend = 60
    val = np.zeros(rend-rstart)
    dispSettings = np.zeros(rend-rstart)
    
    #poses = JT.Testing()
    pose = np.zeros(6)
    pylab.ion()
    pylab.figure()
    pylab.hold(False)

    for i, dispRange in enumerate(range(rstart,rend)):
        dispSettings[i] = dispRange/float(rend)
        drr._volume.SetOpacityLevel(dispSettings[i])
        drr.Render()
        val[i] = reg.GetValue(pose)
        print "Display Setting: %f, Metric Value: %f" % (dispSettings[i], val[i])
        pylab.plot(val)
    
    print "Best setting: ", dispSettings[val.argmin()]
    return val


if __name__ == "__main__":
    Test()